#!/usr/bin/env python3
# -*- coding:utf-8 -*-
# @Author: renjin@bit.edu.cn
# @Date  : 2025-07-08


import threading
import time
import cv2
import os
import json
import argparse
import numpy as np
from typing import Union
from queue import Queue
import platform
from spirems import Publisher, Subscriber, cvimg2sms, sms2cvimg, BaseNode, get_extra_args
import base64


class Colors:
    def __init__(self):
        hexs = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A',
                '92CC17', '3DDB86', '1A9334', '00D4BB', '2C99A8', '00C2FF',
                '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF',
                'FF95C8', 'FF37C7')
        self.palette = [self.hex2rgb(f'#{c}') for c in hexs]
        self.n = len(self.palette)

    def __call__(self, i, bgr=False):
        c = self.palette[int(i) % self.n]
        return (c[2], c[1], c[0]) if bgr else c

    @staticmethod
    def hex2rgb(h):  # rgb order (PIL)
        return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4))


class ObjectSizeStatNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        **kwargs
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(
            self,
            self.__class__.__name__,
            job_name,
            ip=ip,
            port=port,
            param_dict_or_file=param_dict_or_file,
            sms_shutdown=sms_shutdown,
            **kwargs
        )
        self.imshow = self.get_param("imshow", True)
        self.image_topic = self.get_param("image_topic", "")
        self.score_threshold = self.get_param("score_threshold", 0.3)
        self.params_help()

        input_topic = self.image_topic if len(self.image_topic) > 0 else '/' + job_name + '/sensor/image_raw'

        self._image_reader = Subscriber(
            input_topic, 'std_msgs::Null', self.image_callback,
            ip=ip, port=port
        )
        self.colors_obj = Colors()
        self.obj_szs = []

        self.image_queue = Queue()
        self.queue_pool.append(self.image_queue)
        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._image_writer.kill()

    def image_callback(self, msg):
        img = sms2cvimg(msg)
        self.image_queue.put({'msg': msg, 'img': img})

    def run(self):
        while self.is_running():
            img_msg = self.image_queue.get(block=True)
            if img_msg is None:
                break

            img, msg = img_msg['img'], img_msg['msg']
            if 'spirecv_msgs::2DTargets' in msg:
                min_siz = min(msg['spirecv_msgs::2DTargets']['height'], msg['spirecv_msgs::2DTargets']['width'])
                # print(msg['spirecv_msgs::2DTargets']['height'], msg['spirecv_msgs::2DTargets']['width'])
                # print(msg['img_id'])
                for obj in msg['spirecv_msgs::2DTargets']['targets']:
                    obj_sz = obj['bbox'][2] * obj['bbox'][3]
                    self.obj_szs.append(obj_sz)

                if len(self.obj_szs) > 0:
                    data = np.array(self.obj_szs)
                    mean_value = np.mean(data)
                    max_value = np.max(data)
                    min_value = np.min(data)
                    median_value = np.median(data)
                    print("len:", len(self.obj_szs))
                    print("  Mean:   ", mean_value)
                    print("  Maximum:", max_value)
                    print("  Minimum:", min_value)
                    print("  Median: ", median_value)

                if min_siz <= 720:
                    thickness = 1
                elif 720 < min_siz <= 1200:
                    thickness = 2
                else:
                    thickness = 3

                if 'rois' in msg['spirecv_msgs::2DTargets'] and len(msg['spirecv_msgs::2DTargets']['rois']) > 0:
                    roi = msg['spirecv_msgs::2DTargets']['rois'][0]
                    img_roi = img[roi[1]:roi[1] + roi[3], roi[0]:roi[0] + roi[2], :].copy()
                    img = cv2.addWeighted(img, 0.5, np.zeros_like(img, dtype=np.uint8), 0.5, 0)
                    img[roi[1]:roi[1] + roi[3], roi[0]:roi[0] + roi[2], :] = img_roi
                
                masks = []
                result_classid = []
                for obj in msg['spirecv_msgs::2DTargets']['targets']:
                    if 'score' not in obj or obj['score'] >= self.score_threshold:
                        if "segmentation" in obj:
                            obj['segmentation']['counts'] = base64.b64decode(obj['segmentation']['counts'])
                            mask = pycoco_mask.decode(obj['segmentation'])
                            masks.append(mask)
                            result_classid.append(obj['category_id'])

                if len(masks) > 0:
                    alpha = 0.5
                    colors_ = [self.colors_obj(x, True) for x in result_classid]
                    masks = np.asarray(masks, dtype=np.uint8)
                    masks = np.ascontiguousarray(masks.transpose(1, 2, 0))
                    masks = np.asarray(masks, dtype=np.float32)
                    colors_ = np.asarray(colors_, dtype=np.float32)
                    s = masks.sum(2, keepdims=True).clip(0, 1)
                    masks = (masks @ colors_).clip(0, 255)
                    img[:] = masks * alpha + img * (1 - s * alpha)

                for obj in msg['spirecv_msgs::2DTargets']['targets']:
                    if 'tracked_id' in obj:
                        cv2.rectangle(
                            img,
                            (int(obj['bbox'][0]), int(obj['bbox'][1])),
                            (int(obj['bbox'][0] + obj['bbox'][2]), int(obj['bbox'][1] + obj['bbox'][3])),
                            (0, 0, 255),
                            thickness,
                            cv2.LINE_AA
                        )
                        cv2.rectangle(
                            img,
                            (int(obj['bbox'][0]), int(obj['bbox'][1])),
                            (int(obj['bbox'][0] + len(str(obj['tracked_id'])) * 12), int(obj['bbox'][1] + 18)),
                            (0, 0, 0),
                            -1,
                            cv2.LINE_AA
                        )
                        cv2.putText(
                            img,
                            str(obj['tracked_id']),
                            (int(obj['bbox'][0]) + 2, int(obj['bbox'][1]) + 15),
                            cv2.FONT_HERSHEY_SIMPLEX,
                            0.5,
                            (255, 255, 255),
                            1
                        )
                        
                    elif 'score' not in obj or obj['score'] >= self.score_threshold:
                        cv2.rectangle(
                            img,
                            (int(obj['bbox'][0]), int(obj['bbox'][1])),
                            (int(obj['bbox'][0] + obj['bbox'][2]), int(obj['bbox'][1] + obj['bbox'][3])),
                            (0, 0, 255),
                            thickness,
                            cv2.LINE_AA
                        )
                        if obj['bbox'][3] < 50:  # pixel
                            cv2.rectangle(
                                img,
                                (int(obj['bbox'][0]), int(obj['bbox'][1])),
                                (int(obj['bbox'][0] + len(obj['category_name']) * 12), int(obj['bbox'][1] - 18)),
                                (0, 0, 0),
                                -1,
                                cv2.LINE_AA
                            )
                            cv2.putText(
                                img,
                                obj['category_name'],
                                (int(obj['bbox'][0]) + 2, int(obj['bbox'][1]) - 3),
                                cv2.FONT_HERSHEY_SIMPLEX,
                                0.5,
                                (255, 255, 255),
                                1
                            )
                        else:
                            cv2.rectangle(
                                img,
                                (int(obj['bbox'][0]), int(obj['bbox'][1])),
                                (int(obj['bbox'][0] + len(obj['category_name']) * 12), int(obj['bbox'][1] + 18)),
                                (0, 0, 0),
                                -1,
                                cv2.LINE_AA
                            )
                            cv2.putText(
                                img,
                                obj['category_name'],
                                (int(obj['bbox'][0]) + 2, int(obj['bbox'][1]) + 15),
                                cv2.FONT_HERSHEY_SIMPLEX,
                                0.5,
                                (255, 255, 255),
                                1
                            )
                if "fei_cxcy" in msg['spirecv_msgs::2DTargets']:
                    cx = int(msg['spirecv_msgs::2DTargets']['fei_cxcy'][0])
                    cy = int(msg['spirecv_msgs::2DTargets']['fei_cxcy'][1])
                    cv2.circle(img, (cx, cy), 8, (154, 250, 0), 2)

            if self.imshow:
                cv2.imshow('img', img)
                cv2.waitKey(5)

        self.release()
        print('{} quit!'.format(self.__class__.__name__))


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '--config',
        type=str,
        default='default_params.json',
        help='SpireCV2 Config (.json)')
    parser.add_argument(
        '--job-name',
        type=str,
        default='live',
        help='SpireCV Job Name')
    parser.add_argument(
        '--ip',
        type=str,
        default='127.0.0.1',
        help='SpireMS Core IP')
    parser.add_argument(
        '--port',
        type=int,
        default=9094,
        help='SpireMS Core Port')
    # args = parser.parse_args()
    args, unknown_args = parser.parse_known_args()
    if not os.path.isabs(args.config):
        current_path = os.path.abspath(__file__)
        params_dir = os.path.join(current_path[:current_path.find('spirecv-pro') + 11], 'params', 'spirecv2')
        args.config = os.path.join(params_dir, args.config)
    print("--config:", args.config)
    print("--job-name:", args.job_name)
    extra = get_extra_args(unknown_args)

    node = ObjectSizeStatNode(args.job_name, param_dict_or_file=args.config, ip=args.ip, port=args.port, **extra)
    node.join()

